from typing import Union, Dict

import torch


class UnsqueezeDimension:
    """
    Add a new dimension at a specified position for tensors within a dictionary or a standalone tensor.
    """
    def __init__(self, position: int, keys=None) -> None:
        self.position = position
        self.keys = keys

    def _apply(self, tensor: torch.Tensor) -> torch.Tensor:
        """
        Apply unsqueeze operation to the tensor.
        
        :param tensor: Input tensor.
        :return: Tensor with a new dimension added.
        """
        return tensor.unsqueeze(self.position)

    def __call__(self, data: Union[torch.Tensor, Dict[str, Union[torch.Tensor, float]]]) -> Union[torch.Tensor, Dict[str, Union[torch.Tensor, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                data[key] = self._apply(data[key])
        elif isinstance(data, torch.Tensor):
            data = self._apply(data)
        return data


class SqueezeTensor:
    """
    Remove all single-dimensional entries from the shape of tensors, either from specified keys in a dictionary or a standalone tensor.
    """
    def __init__(self, keys=None) -> None:
        self.keys = keys  # No default keys, apply to all tensor keys in a dict if keys is None

    def _apply(self, tensor: torch.Tensor) -> torch.Tensor:
        """
        Apply squeeze operation to the tensor to remove all single dimensions.
        
        :param tensor: Input tensor.
        :return: Tensor after squeezing.
        """
        return tensor.squeeze()

    def __call__(self, data: Union[torch.Tensor, Dict[str, Union[torch.Tensor, float]]]) -> Union[torch.Tensor, Dict[str, Union[torch.Tensor, float]]]:
        if isinstance(data, dict):
            keys_to_process = self.keys if self.keys is not None else data.keys()
            for key in keys_to_process:
                if key in data and isinstance(data[key], torch.Tensor):
                    data[key] = self._apply(data[key])
        elif isinstance(data, torch.Tensor):
            data = self._apply(data)
        return data